#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Datetime: 2020/12/21 21:56
# @Author  : CHEN Wang
# @Site    : 
# @File    : selecting_timing_ability.py
# @Software: PyCharm 

"""
脚本说明：选股择时能力的核心函数
"""

import numpy as np
import pandas as pd
from quant_researcher.quant.project_tool.logger.my_logger import LOG
from quant_researcher.quant.project_tool import hammer
from quant_researcher.quant.project_tool import df_tool
import statsmodels.api as sm


def selecting_timing_ability_model(return_df, risk_free_ret, model_type='tm', calc_return=True):
    """
    选股择时能力的核心函数

    :param pd.DataFrame return_df: 需要计算的标的的收益率和对应基准的收益率序列，
                                   第一列是日期，第二列是标的，第三列是基准，列名分别是标的和基准的代码
                    +------------+---------+---------+
                    | trade_date |  110022 |  000300 |
                    +------------+---------+---------+
                    | 2020-07-13 |  0.0406 |  0.021  |
                    | 2020-07-14 |  -0.002 | -0.0095 |
                    | 2020-07-15 |  0.0218 | -0.0129 |
                    | 2020-07-16 | -0.0703 | -0.0481 |
                    | 2020-07-17 |  0.015  |  0.0063 |
                    +------------+---------+---------+
    :param float risk_free_ret: 对应的无风险收益率
    :param str model_type: 计算模型，支持tm、hm、cl
    :param bool calc_return: 是否计算收益，默认计算
    :return: dict
                +--------------------------+------------+
                |            键            |     值     |
                +--------------------------+------------+
                |        asset_code        |   110022   |
                |      benchmark_code      |   000300   |
                |     stock_ability_tm     |  0.005937  |
                | stock_ability_pvalue_tm  |  0.067825  |
                |    timing_ability_tm     | -14.220607 |
                | timing_ability_pvalue_tm |  0.014965  |
                |        total_ret         |  0.01334   |
                |       stock_ret_tm       |  0.136555  |
                |      timing_ret_tm       | -0.098852  |
                |      market_ret_tm       | -0.016209  |
                +--------------------------+------------+
    """
    asset_code = return_df.columns[1]
    benchmark_code = return_df.columns[2]
    data = {}
    # 不改变传进来的return_df
    tmp_df = return_df.copy()
    tmp_df['y'] = tmp_df[asset_code] - risk_free_ret
    tmp_df['x1'] = tmp_df[benchmark_code] - risk_free_ret
    tmp_df['x2'] = tmp_df[benchmark_code] - risk_free_ret
    if model_type == 'tm':
        tmp_df.loc[:, 'x2'] = tmp_df['x2'] * (tmp_df[benchmark_code] - risk_free_ret)
    if model_type == 'hm':
        tmp_df.loc[:, 'x2'] = tmp_df['x2'].apply(lambda x: max(x, 0))
    if model_type == 'cl':
        tmp_df.loc[:, 'x1'] = tmp_df['x1'].apply(lambda x: min(x, 0))
        tmp_df.loc[:, 'x2'] = tmp_df['x2'].apply(lambda x: max(x, 0))
    del tmp_df[benchmark_code]
    del tmp_df[asset_code]
    hammer.slim(tmp_df, a_f=1, name=f'{asset_code}结果', deal_float=False)
    LOG.info(f'{asset_code}结果行数={tmp_df.shape[0]}')
    if tmp_df.shape[0] < 3:
        LOG.warning('statsmodels 的 omni_normtest 要求至少 8 个数据')
        return

    x_s = df_tool.df_to_np_array(tmp_df[['x1', 'x2']])
    y_s = df_tool.df_to_np_array(tmp_df[['y']])

    x_s = np.c_[np.ones((tmp_df.shape[0], 1)), x_s]  # 加上一个常数项在最前
    y_s = np.array(y_s, dtype='float')

    model = sm.OLS(y_s, exog=x_s)
    fit_ans = model.fit()

    # print(fit_ans.summary())

    w_ans = fit_ans.params
    p_ans = fit_ans.pvalues

    if model_type == 'tm' or model_type == 'hm':
        timing_ability = round(float(w_ans[2]), 6)
        timing_ability_p_value = round(float(p_ans[2]), 6)
    else:
        # Chang-Lewellen时，用 beta_2 - beta_1，那这时的p-value怎么计算呢？
        timing_ability = round(float(w_ans[2] - w_ans[1]), 6)
        timing_ability_p_value = round(float(max(p_ans[2], p_ans[1])), 6)
    data.update({
        f'stock_ability_{model_type}': round(float(w_ans[0]), 6),
        f'stock_ability_pvalue_{model_type}': round(float(p_ans[0]), 6),
        f'timing_ability_{model_type}': timing_ability,
        f'timing_ability_pvalue_{model_type}': timing_ability_p_value
    })

    if calc_return:
        stock_ret = timing_ret = market_ret = np.nan
        total_ret = (tmp_df['y'] + 1).prod() - 1
        if model_type == 'tm' or model_type == 'hm':
            stock_ret = tmp_df.shape[0] * w_ans[0]
            timing_ret = (timing_ability * tmp_df['x2'] + 1).prod() - 1
            market_ret = (w_ans[1] * tmp_df['x1'] + 1).prod() - 1
        data.update({
            f'total_ret': round(float(total_ret), 6),
            f'stock_ret_{model_type}': round(float(stock_ret), 6),
            f'timing_ret_{model_type}': round(float(timing_ret), 6),
            f'market_ret_{model_type}': round(float(market_ret), 6)
        })

    # import pandas as pd
    # hammer.slim(pd.DataFrame([data]), a_f=0, show_index=1, show_all=1)

    return data


if __name__ == '__main__':
    from quant_researcher.quant.datasource_fetch.fund_api.fund_nav_related import get_fund_return
    from quant_researcher.quant.datasource_fetch.index_api.index_price_related import \
        get_index_return

    fund_code = '110022'
    index_code = '000300'
    start_date = '2020-01-01'
    end_date = '2020-12-01'

    fund_df = get_fund_return(fund_code=fund_code, start_date=start_date, end_date=end_date)
    index_df = get_index_return(index_code=index_code, start_date=start_date, end_date=end_date)
    fund_df = fund_df.rename(columns={'daily_return': '110022'})
    index_df = index_df.rename(columns={'daily_return': '000300'})
    return_df = fund_df.merge(index_df, how='inner', on='end_date')
    return_df = return_df[['end_date', '110022', '000300']]
    aaa = selecting_timing_ability_model(return_df, 0.015, 'tm')
